import datetime
import time
import pymysql.cursors
from fastapi import APIRouter, HTTPException, Query, Request, WebSocket
import requests
import json
from functools import lru_cache
from config import get_settings
from sse_manager import SSEManager

# Connect to the database
connection = pymysql.connect(host='localhost',
                             user='root',
                             password='pp2bgL@L',
                             db='history',
                             charset='utf8mb4',
                             cursorclass=pymysql.cursors.DictCursor)


# 保持数据库连接
def getconnection():
    connection.ping(reconnect=True)
    return connection


api_wenxinyiyan = APIRouter()
api_wenxinyige = APIRouter()
sse_manager = SSEManager()


# 使用lru_cache缓存access_token
@lru_cache(maxsize=1)
def get_access_token(service):
    settings = get_settings()
    if service == "yan":
        url = f"{settings.baidu_token_url}?grant_type=client_credentials&client_id={settings.baidu_client_id}&client_secret={settings.baidu_client_secret}"
    if service == "ge":
        url = f"{settings.baidu_token_url}?grant_type=client_credentials&client_id={settings.baidu_draw_id}&client_secret={settings.baidu_draw_secret}"
    response = requests.post(url)
    if response.status_code != 200:
        raise HTTPException(status_code=500, detail="Failed to retrieve access_token")

    token_data = response.json()
    print(token_data)
    return token_data.get("access_token")


@api_wenxinyiyan.get("/chat")
async def chat(content: str = Query(..., description="聊天内容")):
    settings = get_settings()
    if "请为我画" in content:
        access_token = get_access_token("ge")
        print(access_token)
        print("调用文心一格")
        url = f"{settings.baidu_draw_url}?access_token={access_token}"

        payload = json.dumps({
            "text": content,
            "resolution": "1024*1024",
            "style": "写实风格"
        })
        headers = {
            "Accept": "application/json",
            "Content-Type": "application/json"
        }
        response = requests.post(url, headers=headers, data=payload)

        if response.status_code != 200:
            raise HTTPException(status_code=500, detail="与文心一格API通信失败")

        chat_response = response.json()
        print(chat_response)
        # 从响应中获取taskId
        task_id = chat_response.get("data", {}).get("taskId")
        print(task_id)
        task_url = f"{settings.baidu_draw_result_url}?access_token={access_token}"
        payload = json.dumps({
            "taskId": task_id
        })
        headers = {
            'Content-Type': 'application/json',
            'Accept': 'application/json'
        }

        image_url = ""
        while not image_url:
            response1 = requests.request("POST", task_url, headers=headers, data=payload)
            image_response = response1.json()
            image_url = image_response.get("data", {}).get("img", "")
            if not image_url:
                time.sleep(2)  # 等待一段时间后重试

        # 将图片的URL发送给前端
        await sse_manager.publish(image_url)
        print(image_url)
        now = datetime.datetime.now()
        time1 = now.strftime('%Y-%m-%d %H:%M:%S')
        try:
            with getconnection().cursor() as cursor:
                sql = "INSERT INTO `communication` (`type`,`communications`,`time`) VALUES (%s, %s, %s)"
                cursor.execute(sql, ("USER", content, time1))
                connection.commit()
                sql = "INSERT INTO `communication` (`type`,`communications`,`time`) VALUES (%s, %s, %s)"
                cursor.execute(sql, ("AI", image_url, time1))
                connection.commit()
                cursor.close()
        except Exception as e:
            cursor.close()
        result = chat_response.get("result", "")
        # 返回图片的URL作为响应
        return {"response": image_url}
    else:
        access_token = get_access_token("yan")
        print(access_token)
        url = f"{settings.baidu_chat_url}?access_token={access_token}"
        cursor1 = getconnection().cursor()
        sql2 = "SELECT * FROM `chat` "
        sql1 = "INSERT INTO `chat` (`role`,`content`) VALUES (%s, %s)"
        cursor1.execute(sql1, ("user", content))
        cursor1.execute(sql2)
        tables = cursor1.fetchall()
        connection.commit()
        results_list = list(tables)
        print(results_list)
        payload = json.dumps({
            "messages": results_list
        })
        headers = {
            'Content-Type': 'application/json'
        }

        response = requests.post(url, headers=headers, data=payload)

        if response.status_code != 200:
            raise HTTPException(status_code=500, detail="与文心一言API通信失败")

        chat_response = response.json()
        print(chat_response)
        now = datetime.datetime.now()
        time1 = now.strftime('%Y-%m-%d %H:%M:%S')
        try:
            with getconnection().cursor() as cursor:
                sql = "INSERT INTO `communication` (`type`,`communications`,`time`) VALUES (%s, %s, %s)"
                cursor.execute(sql, ("USER", content, time1))
                connection.commit()
                sql = "INSERT INTO `communication` (`type`,`communications`,`time`) VALUES (%s, %s, %s)"
                sql1 = "INSERT INTO `chat` (`role`,`content`) VALUES (%s, %s)"
                cursor.execute(sql, ("AI", chat_response.get("result", ""), time1))
                cursor.execute(sql1, ("assistant", chat_response.get("result", "")))
                connection.commit()
                cursor.close()
        except Exception as e:
            cursor.close()
        result = chat_response.get("result", "")

        # 将结果分割成多个部分,每个部分发送一次 SSE 事件
        for i in range(0, len(result), 1):
            chunk = result[i:i + 1]
            await sse_manager.publish(chunk)
            print(chunk, end="")  # 在控制台输出对话记录

        print("\n")  # 在控制台添加换行符
        return {"response": result}


@api_wenxinyiyan.get("/showhistory")
async def showhistory():
    try:
        with getconnection().cursor() as cursor:
            sql = "SELECT * FROM `communication`"
            cursor.execute(sql)
            result = cursor.fetchall()
            print(11111111111)
            cursor.close()
            return {"response": result}
    except Exception as e:
        cursor.close()


@api_wenxinyiyan.get("/events")
async def subscribe_events(request: Request):
    return await sse_manager.subscribe(request)